# -*- coding: utf-8 -*-
import contextlib
import datetime

from sqlalchemy import exc as sql_exc
from sqlalchemy import pool
from sqlalchemy.engine import create_engine
from sqlalchemy.orm import scoped_session
from sqlalchemy.orm import sessionmaker

from multidelivery import config
from multidelivery.db.sqlalchemy import models


class DBSessionManager(object):
    def __init__(self):
        self.session_map = {}
        self.init_db()

    def init_db(self):
        self.session_map = {}
        self.create_sessions()

    def create_sessions(self):
        self.session_map['master'] = self.create_single_session(config.DB_URI)
        if config.SLAVE_DB_URI:
            self.session_map['slave'] = self.create_single_session(config.SLAVE_DB_URI)

    @classmethod
    def create_single_session(cls, url, scopefunc=None):
        engine_opts = {}
        engine_opts.update(config.SQLALCHEMY_ENGINE_OPTIONS)
        engine_opts['poolclass'] = pool.NullPool
        engine = create_engine(url, **config.SQLALCHEMY_ENGINE_OPTIONS)
        return scoped_session(sessionmaker(expire_on_commit=False, bind=engine), scopefunc=scopefunc)

    def get_session(self, slave=False):
        try:
            if slave and config.SLAVE_DB_URI:
                return self.session_map['slave']
            else:
                return self.session_map['master']
        except KeyError:
            raise KeyError('{} not created, check your DB_SETTINGS'.format(slave))
        except IndexError:
            raise IndexError('cannot get names from DB_SETTINGS')

    @contextlib.contextmanager
    def session_ctx(self, slave=False):
        DBSession = self.get_session(slave)
        session = DBSession()
        try:
            yield session
            session.commit()
        except Exception:
            session.rollback()
            raise
        finally:
            session.expunge_all()
            session.close()


db_manager = DBSessionManager()


def create_worker(hostname):
    worker_dict = {
        "host": hostname
    }

    work_node = models.Worker(**worker_dict)
    with db_manager.session_ctx() as session:
        exist = session.query(models.Worker.id).filter_by(host=hostname).scalar() is not None
        if not exist:
            session.add(work_node)
        else:
            update_worker_heartbeat(hostname)


def update_worker_heartbeat(hostname):
    with db_manager.session_ctx() as session:
        session.query(models.Worker).filter_by(host=hostname).update({"updated_at": datetime.datetime.now()})
